package com.boulderai.mcp.websocket;

import com.boulderai.mcp.service.McpService;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.HashMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

/**
 * MCP WebSocket处理器
 * 处理WebSocket连接和消息
 */
@Component
public class McpWebSocketHandler extends TextWebSocketHandler {

    @Autowired
    private McpService mcpService;

    private final ObjectMapper objectMapper = new ObjectMapper();
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.put(session.getId(), session);
        System.out.println("WebSocket connection established: " + session.getId());
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        try {
            String payload = message.getPayload();
            System.out.println("Received WebSocket message: " + payload);
            
            JsonNode messageNode = objectMapper.readTree(payload);
            JsonNode typeNode = messageNode.get("type");
            
            if (typeNode == null) {
                sendError(session, "Message type is missing");
                return;
            }
            
            String type = typeNode.asText();
            
            switch (type) {
                case "connect":
                    handleConnect(session, messageNode);
                    break;
                case "disconnect":
                    handleDisconnect(session, messageNode);
                    break;
                case "list_tools":
                    handleListTools(session, messageNode);
                    break;
                case "call_tool":
                    handleCallTool(session, messageNode);
                    break;
                case "list_resources":
                    handleListResources(session, messageNode);
                    break;
                case "read_resource":
                    handleReadResource(session, messageNode);
                    break;
                default:
                    sendError(session, "Unknown message type: " + type);
            }
        } catch (Exception e) {
            System.err.println("Error processing WebSocket message: " + e.getMessage());
            e.printStackTrace();
            sendError(session, "Error processing message: " + e.getMessage());
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session.getId());
        System.out.println("WebSocket connection closed: " + session.getId());
    }

    private void handleConnect(WebSocketSession session, JsonNode message) {
        String serverUri = message.get("serverUri").asText();
        // 检查sessionId是否存在
        if (!message.has("sessionId") || message.get("sessionId").asText().isEmpty()) {
            sendError(session, "sessionId is required for connect operation");
            return;
        }
        String sessionId = message.get("sessionId").asText();
        
        mcpService.connect(sessionId, serverUri)
                .thenRun(() -> sendConnectResponse(session, sessionId, "Connected successfully"))
                .exceptionally(ex -> {
                    sendError(session, "Connection failed: " + ex.getMessage());
                    return null;
                });
    }

    private void handleDisconnect(WebSocketSession session, JsonNode message) {
        // 检查sessionId是否存在
        if (!message.has("sessionId") || message.get("sessionId").asText().isEmpty()) {
            sendError(session, "sessionId is required for disconnect operation");
            return;
        }
        String sessionId = message.get("sessionId").asText();
        
        mcpService.disconnect(sessionId)
                .thenRun(() -> sendResponse(session, "disconnect", "success", "Disconnected successfully"))
                .exceptionally(ex -> {
                    sendError(session, "Disconnect failed: " + ex.getMessage());
                    return null;
                });
    }

    private void handleListTools(WebSocketSession session, JsonNode message) {
        // 检查sessionId是否存在
        if (!message.has("sessionId") || message.get("sessionId").asText().isEmpty()) {
            sendError(session, "sessionId is required for list_tools operation");
            return;
        }
        String sessionId = message.get("sessionId").asText();
        
        mcpService.listTools(sessionId)
                .thenAccept(tools -> sendResponse(session, "list_tools", "success", tools))
                .exceptionally(ex -> {
                    sendError(session, "Failed to list tools: " + ex.getMessage());
                    return null;
                });
    }

    private void handleCallTool(WebSocketSession session, JsonNode message) {
        // 检查sessionId是否存在
        if (!message.has("sessionId") || message.get("sessionId").asText().isEmpty()) {
            sendError(session, "sessionId is required for call_tool operation");
            return;
        }
        String sessionId = message.get("sessionId").asText();
        String toolName = message.get("toolName").asText();
        JsonNode arguments = message.get("arguments");
        
        mcpService.callTool(sessionId, toolName, arguments)
                .thenAccept(result -> sendResponse(session, "call_tool", "success", result))
                .exceptionally(ex -> {
                    sendError(session, "Tool call failed: " + ex.getMessage());
                    return null;
                });
    }

    private void handleListResources(WebSocketSession session, JsonNode message) {
        // 检查sessionId是否存在
        if (!message.has("sessionId") || message.get("sessionId").asText().isEmpty()) {
            sendError(session, "sessionId is required for list_resources operation");
            return;
        }
        String sessionId = message.get("sessionId").asText();
        
        mcpService.listResources(sessionId)
                .thenAccept(resources -> sendResponse(session, "list_resources", "success", resources))
                .exceptionally(ex -> {
                    sendError(session, "Failed to list resources: " + ex.getMessage());
                    return null;
                });
    }

    private void handleReadResource(WebSocketSession session, JsonNode message) {
        // 检查sessionId是否存在
        if (!message.has("sessionId") || message.get("sessionId").asText().isEmpty()) {
            sendError(session, "sessionId is required for read_resource operation");
            return;
        }
        String sessionId = message.get("sessionId").asText();
        String resourceUri = message.get("resourceUri").asText();
        
        mcpService.readResource(sessionId, resourceUri)
                .thenAccept(content -> sendResponse(session, "read_resource", "success", content))
                .exceptionally(ex -> {
                    sendError(session, "Failed to read resource: " + ex.getMessage());
                    return null;
                });
    }

    private void sendResponse(WebSocketSession session, String type, String status, Object data) {
        try {
            Map<String, Object> response = new HashMap<>();
            response.put("type", type);
            response.put("status", status);
            response.put("data", data);
            String responseJson = objectMapper.writeValueAsString(response);
            session.sendMessage(new TextMessage(responseJson));
        } catch (IOException e) {
            System.err.println("Failed to send response: " + e.getMessage());
        }
    }

    private void sendConnectResponse(WebSocketSession session, String sessionId, String message) {
        try {
            Map<String, Object> response = new HashMap<>();
            response.put("type", "connect");
            response.put("status", "success");
            response.put("sessionId", sessionId);
            response.put("data", message);
            String responseJson = objectMapper.writeValueAsString(response);
            session.sendMessage(new TextMessage(responseJson));
        } catch (IOException e) {
            System.err.println("Failed to send connect response: " + e.getMessage());
        }
    }

    private void sendError(WebSocketSession session, String error) {
        try {
            Map<String, Object> response = new HashMap<>();
            response.put("type", "error");
            response.put("status", "error");
            response.put("message", error);
            String responseJson = objectMapper.writeValueAsString(response);
            session.sendMessage(new TextMessage(responseJson));
        } catch (IOException e) {
            System.err.println("Failed to send error: " + e.getMessage());
        }
    }
}